"""
    @brief:Particle filter algorithm
    @Editor:CJH
    @Date:2025/4/24
"""
import threading

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
import GPy

import cv2

from utils import Map_init2


# 真实气体源位置（用于仿真）


class ParticleFilter:
    def __init__(self, particle_nums, canvas_size=(900, 540), split_x_nums=10, split_y_nums=15):
        self.particle_history = None
        self.Particle_nums = particle_nums  # 无需过多的粒子，填满每一个栅格即可，加速收敛
        # 构建一些机器人
        Bot_elements = {"Bot1": Map_init2.BOT("Bot1", (0, 0), 20)}
        self.terminal = Map_init2.Terminal(Bot_elements, canvas_size=canvas_size, split_x_nums=split_x_nums,
                                           split_y_nums=split_y_nums)
        self.terminal.render()
        self.H, self.W = canvas_size
        self.N = 5000  # 粒子数量
        self.sigma = 70  # 扩散速度参数
        self.sigma_m = 0.7  # 测量噪声标准差
        self.gp_model = None
        self.training_data = []  # 存储训练数据 (位置, 浓度)
        self.max_training_points = 100  # 限制训练数据量，避免计算过慢

    # 初始化粒子
    def initialize_particles(self):
        particles = np.random.randint(0, [self.W, self.H], size=(self.N, 2))  # 随机分布
        weights = np.ones(self.N) / self.N  # 均匀权重
        return particles, weights

    def update_gp_model(self, robot_pos, C_measured):
        """更新高斯过程模型"""
        # 添加新数据
        self.training_data.append([robot_pos[0], robot_pos[1], C_measured])
        if len(self.training_data) > self.max_training_points:
            self.training_data.pop(0)  # 移除最旧的数据

        # 提取训练数据
        data = np.array(self.training_data)
        X_train = data[:, :2]  # 位置 (x, y)
        y_train = data[:, 2]  # 浓度

        # 初始化或更新GP模型
        kernel = GPy.kern.RBF(input_dim=2, variance=1.0, lengthscale=50.0)
        if self.gp_model is None:
            self.gp_model = GPy.models.GPRegression(X_train, y_train[:, None], kernel, noise_var=self.sigma_m ** 2)
        else:
            self.gp_model.set_XY(X_train, y_train[:, None])
        self.gp_model.optimize()

    def likelihood_with_gp(self, robot_pos):
        """使用高斯过程构建似然函数"""
        # 假设粒子位置为气体源，用GP预测机器人位置的浓度
        robot_pos_array = np.array([[robot_pos[0], robot_pos[1]]])
        C_pred, C_var = self.gp_model.predict(robot_pos_array)
        C_pred = C_pred[0, 0]
        sigma_pred = np.sqrt(C_var[0, 0])

        # 测量值
        C_measured = self.training_data[-1][2]  # 最新的测量值
        # 高斯似然
        likelihood = np.exp(-((C_measured - C_pred) ** 2) / (2 * sigma_pred ** 2)) / np.sqrt(
            2 * np.pi * sigma_pred ** 2)

        return likelihood

    def is_in_non_target(self, x, y, visited_non_target):
        """检查粒子 (x, y) 是否在 visited_non_target 的任一矩形区域内"""
        for rect in visited_non_target:
            x_min, y_min, x_max, y_max = rect
            if x_min <= x <= x_max and y_min <= y <= y_max:
                return True
        return False

    # 重采样
    def resample(self, particles, weights):
        # visited_non_target = self.mark_non_source()
        indices = np.random.choice(np.arange(self.N), size=self.N, p=weights)
        particles = particles[indices]
        # 添加扰动
        # 移除非目标位置的粒子
        for i in range(self.N):
            x, y = particles[i]
            if np.random.rand() < 0.1:  # 10% 的粒子添加扰动
                dx, dy = np.random.normal(0, 5, 2)  # 扰动范围可调整
                new_x = max(0, min(self.W - 1, int(x + dx)))
                new_y = max(0, min(self.H - 1, int(y + dy)))
                particles[i] = (new_x, new_y)

        weights = np.ones(self.N) / self.N
        return particles, weights

    # def mark_non_source(self):
    #     # 标记非气体源栅格
    #     visited_non_source = []
    #     max_mass_thresh = np.max(self.terminal.mass_fraction_map)  # 最大的质量分数，相当于最低的测量阻值
    #     quality_scores_map = self.terminal.grid_map * self.terminal.mass_fraction_map
    #     indexes = np.array(
    #         np.where(np.bitwise_and(quality_scores_map > 0, quality_scores_map < max_mass_thresh * 0.9)))[::-1].T
    #     # 将grid转换xyxy的矩形
    #     for index in indexes:
    #         visited_non_source.append(
    #             [*(index * self.terminal.box_min_step), *(index + np.array([1, 1])) * self.terminal.box_min_step])
    #     return visited_non_source

    def get_cluster_centers(self, particles, num_clusters):
        """使用 KMeans 聚类获取粒子分布的中心"""
        if len(particles) < num_clusters:
            return []
        kmeans = KMeans(n_clusters=num_clusters)
        kmeans.fit(particles)
        return kmeans.cluster_centers_

    def bot_controller(self):
        # 初始化，粒子环境
        particles, weights = self.initialize_particles()
        self.particle_history = [particles.copy()]

        while True:
            self.terminal.draw_elements_on_canvas(csv_path=r'D:\Gas_detector\data\gas_mass_data\export5.csv')
            # plot
            cv2.imshow("mass", cv2.resize(self.terminal.mass_fraction_map, (600, 900)))
            key = cv2.waitKey(1)
            # 机器人仿真控制
            for bot in self.terminal.element_collection.values():
                if key == 119:
                    self.terminal.take_actions("upper", False, f"{bot.name}")
                elif key == 115:
                    self.terminal.take_actions("lower", False, f"{bot.name}")
                elif key == 100:
                    self.terminal.take_actions("right", False, f"{bot.name}")
                elif key == 97:
                    self.terminal.take_actions("left", False, f"{bot.name}")
                elif key == 113:
                    self.terminal.reset()
                else:
                    continue
                # 可视化机器人的位置
                # print(self.terminal.mass_fraction_map[bot.current_grid_pos[1],bot.current_grid_pos[0]])
                particles, weights = self.particle_filter(bot, particles, weights)
                # 使用聚类k-means，可视化粒子簇中心
                self.particle_history.append(particles)
                self.terminal.draw_elements_on_canvas(csv_path=r'D:\Gas_detector\data\gas_mass_data\export5.csv')
                self.visualize()

    # 粒子滤波主函数
    def particle_filter(self, bot, particles, weights):
        z_t = self.terminal.mass_fraction_map[
                  bot.current_grid_pos[1], bot.current_grid_pos[0]] * np.exp(np.random.normal(0, self.sigma_m))
        """更新高斯过程模型"""
        self.update_gp_model(bot.current_pixel_pos, z_t)
        print(z_t)
        """根据GP似然更新粒子权重"""
        if self.gp_model is None:
            return  # 如果GP模型未初始化，跳过
        # 通过似然函数更新权重

        for i in range(self.N):
            # print(i)
            weights[i] *= self.likelihood_with_gp(bot.current_pixel_pos)
            # print(weights[i])
        # 归一化权重
        weights += 1e-10
        weights /= np.sum(weights)
        particles, weights = self.resample(particles, weights)
        max_weight_index = np.argmax(weights)
        min_weight_index=np.argmin(weights)
        print(min_weight_index)
        return particles, weights

    def visualize(self):
        plt.figure(figsize=(10, 8))
        # 绘制地图和真实源
        plt.imshow(np.zeros((self.H, self.W)), cmap='gray', origin='lower')
        known_max_pos = np.unravel_index(np.argmax(self.terminal.mass_fraction_map),
                                         self.terminal.mass_fraction_map.shape)[::-1] * self.terminal.box_min_step

        plt.scatter(known_max_pos[0], known_max_pos[1], c='red', marker='x', s=200, label='True Source')
        # 可视化机器人的位置
        for bot in self.terminal.element_collection.values():
            plt.scatter(bot.x, bot.y, c='green', marker='o', s=200, label=f'BOT-{bot.name} Source')
        # 绘制初始和最终粒子
        initial_particles = self.particle_history[0]
        final_particles = self.particle_history[-1]
        clusters = np.array(self.get_cluster_centers(final_particles, 5), np.int32)
        for peak in clusters:
            plt.scatter(peak[0], peak[1], c='yellow', marker='x', s=250)
        clusters = np.array(self.get_cluster_centers(clusters, 1), np.int32)
        for peak in clusters:
            plt.scatter(peak[0], peak[1], c='white', marker='x', s=700, label='visual Source')

        plt.scatter(initial_particles[:, 0], initial_particles[:, 1], c='blue', alpha=0.3, s=10,
                    label='Initial Particles')
        plt.scatter(final_particles[:, 0], final_particles[:, 1], c='green', alpha=0.5, s=10,
                    label='Final Particles')

        plt.title('Particle Filter with Distance Decay Model')
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.gca().invert_yaxis()  # 反转 y 轴
        plt.grid(True)
        plt.legend()
        plt.show()
        # plt.savefig('particle_filter_result.png')


if __name__ == "__main__":
    PF = ParticleFilter(particle_nums=1000)
    # particle_history = particle_filter()
    PF.bot_controller()
    # PF.visualize(particle_history)
